package org.yinweichen.llm_dataset_backend.service.impl;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.RequiredArgsConstructor;
import org.springframework.http.ResponseEntity;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Propagation;
import org.springframework.transaction.annotation.Transactional;
import org.springframework.transaction.support.TransactionSynchronizationManager;
import org.yinweichen.llm_dataset_backend.DTO.BatchGenerateRequest;
import org.yinweichen.llm_dataset_backend.entity.*;
import org.yinweichen.llm_dataset_backend.repository.*;
import org.yinweichen.llm_dataset_backend.service.EvalService;
import org.yinweichen.llm_dataset_backend.service.PointsService;
import org.yinweichen.llm_dataset_backend.util.ModelApiClient;

import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

@Service
@RequiredArgsConstructor
@Transactional
public class EvalServiceImpl implements EvalService {
    private final EvalTaskRepository evalTaskRepository;
    private final StandardQuestionRepository standardQuestionRepository;
    private final StandardAnswerRepository standardAnswerRepository;
    private final LlmAnswerRepository llmAnswerRepository;
    private final EvaluateRepository evaluateRepository;
    private final PointsService pointsService;

    @Override
    @Transactional(rollbackFor = Exception.class, propagation = Propagation.REQUIRED)
    public ResponseEntity<?> runEvaluation(Long taskId) {
        ModelApiClient.setAPIKey();
        System.out.println("开始执行评测任务: " + taskId);
        System.out.println("当前事务状态: " + (TransactionSynchronizationManager.isActualTransactionActive() ? "活跃" : "非活跃"));
        try {
            EvalTask evalTask = evalTaskRepository.findById(taskId)
                    .orElseThrow(() -> new RuntimeException("未找到评测任务"));

            // 通过ID直接更新任务为运行中状态
            evalTaskRepository.updateStatusById(taskId, "RUNNING");

            // 1. 生成模型回答
            BatchGenerateRequest generateRequest = new BatchGenerateRequest(
                    evalTask.getModels(),
                    evalTask.getQuestions(),
                    evalTask.getVersion()
            );

            // 分批处理模型回答生成
            for (int i = 0; i < 2; i++) { // 最多重试2次
                try {
                    batchGetModelAnswers(generateRequest, taskId);
                    break;
                } catch (Exception e) {
                    if (i == 1) throw e;
                    Thread.sleep(5000); // 等待5秒后重试
                }
            }

            // 2. 执行评测
            ResponseEntity<?> evalResponse = batchEvaluateAnswers(taskId);

            if (!evalResponse.getStatusCode().is2xxSuccessful()) {
                throw new RuntimeException(evalResponse.getBody().toString());
            }

            // 通过ID直接更新任务为已完成状态
            evalTaskRepository.updateStatusById(taskId, "COMPLETED");

            System.out.println("评测任务执行成功: " + taskId);
            return ResponseEntity.ok(Map.of(
                "success", true,
                "message", "评测任务执行完成"
            ));
        } catch (Exception e) {
            System.err.println("评测任务执行失败: " + taskId);
            throw new RuntimeException("评测任务执行失败: "+e.getMessage(), e);
        }
    }

    @Transactional(rollbackFor = {Exception.class, RuntimeException.class}, propagation = Propagation.REQUIRES_NEW)
    public List<LlmAnswer> batchGetModelAnswers(BatchGenerateRequest request, Long taskId) {
        System.out.println("开始生成模型回答, 请求参数: "+request.toString());
        // 参数校验
        if (request.questionIds() == null || request.questionIds().isEmpty()) {
            throw new IllegalArgumentException("Question IDs cannot be null or empty");
        }
        if (request.modelNames() == null || request.modelNames().isEmpty()) {
            throw new IllegalArgumentException("Model names cannot be null or empty");
        }
        EvalTask task = evalTaskRepository.findById(taskId)
                .orElseThrow(() -> new RuntimeException("未找到评测任务"));
        // 获取所有标准问题
        List<StandardQuestion> questions = new ArrayList<>();
        String version=request.version();
        for(Long questionId : request.questionIds()) {
            questions.add(standardQuestionRepository.findById(new StandardQuestionId(questionId,version))
                    .orElseThrow(() -> new RuntimeException("Question not found with id: " + questionId)));
        }
        System.out.println("问题数量: "+questions.size());

        // 分批处理，每批5个模型
        int batchSize = 5;
        List<LlmAnswer> allAnswers = new ArrayList<>();
        String prompt="You are a helpful assistant in Computer Science, please answer the following question within 400 words.";

        for (int i = 0; i < request.modelNames().size(); i += batchSize) {
            System.out.println("生成中...");
            List<String> batchModels = request.modelNames().subList(i, Math.min(i + batchSize, request.modelNames().size()));
            List<LlmAnswer> batchAnswers = batchModels.parallelStream()
                .flatMap(modelName -> questions.stream()
                    .map(question -> {
                        // 调用ModelApiClient获取AI回答
                        ModelApiClient.setCurrModel(modelName);
                        String aiAnswer = ModelApiClient.getAnswer(prompt,question.getQuestion());

                        // 创建回答记录
                        LlmAnswer answer = new LlmAnswer();
                        answer.setModel(modelName);
                        answer.setQuestion(question);
                        answer.setAnswer(aiAnswer);

                        // 获取完整的EvalTask对象
                        answer.setTask(task);

                        return answer;
                    }))
                .collect(Collectors.toList());

            // 每批处理完后立即提交事务
            List<LlmAnswer> savedAnswers = llmAnswerRepository.saveAll(batchAnswers);
            allAnswers.addAll(savedAnswers);
            System.out.println("已处理模型批次: " + (i/batchSize + 1) + "/" + (int)Math.ceil((double)request.modelNames().size()/batchSize));
        }

        return allAnswers;
    }

    @Transactional(rollbackFor = {Exception.class, RuntimeException.class,IllegalArgumentException.class}, propagation = Propagation.REQUIRES_NEW)
    public ResponseEntity<?> batchEvaluateAnswers(Long taskId) {
        try {
            //读取评测任务涉及的所有大模型回答
            List<LlmAnswer> llmAnswers = llmAnswerRepository.findAllByTaskId(taskId);
            EvalTask evalTask = evalTaskRepository.findById(taskId)
                    .orElseThrow(() -> new RuntimeException("未找到评测任务"));
            // 记录任务信息
            for (LlmAnswer llmAnswer : llmAnswers) {
                StandardQuestion question = llmAnswer.getQuestion();
                System.out.println("处理问题: " + question.getId() + ", 模型回答ID: " + llmAnswer.getId());
            }
            System.out.println("开始评测任务: " + evalTask.getTitle());
            System.out.println("任务描述: " + evalTask.getDescription());

            // 获取并验证评测配置
            String prompt = evalTask.getPrompt();
            if (prompt == null || prompt.trim().isEmpty()) {
                throw new IllegalArgumentException("评测提示词(prompt)不能为空");
            }
            if (prompt.length() > 1000) {
                throw new IllegalArgumentException("评测提示词(prompt)长度不能超过1000字符");
            }

            // 记录脱敏后的prompt(前20字符+...)
            String maskedPrompt = prompt.length() > 20 ?
                prompt.substring(0, 20) + "..." : prompt;
            System.out.println("使用评测提示词: " + maskedPrompt);

            String version = evalTask.getVersion();
            List<String> judges = evalTask.getJudges();

            // 分批处理，每批10个答案
            int batchSize = 10;

            for(String judge : judges) {
                ModelApiClient.setCurrModel(judge);
                System.out.println("开始使用裁判模型: " + judge);

                for (int i = 0; i < llmAnswers.size(); i += batchSize) {
                    List<LlmAnswer> batchAnswers = llmAnswers.subList(i, Math.min(i + batchSize, llmAnswers.size()));
                    List<Evaluate> batchEvaluates = new ArrayList<>();

                    for (LlmAnswer llmAnswer : batchAnswers) {
                        StandardQuestion standardQuestion = llmAnswer.getQuestion();
                        StandardAnswer standardAnswer = standardAnswerRepository.findById(new StandardAnswerId(standardQuestion.getId(), version))
                                .orElseThrow(() -> new RuntimeException("未找到对应的StandardAnswer"));

                        String jsonData = generateJsonData(standardAnswer, standardQuestion, llmAnswer);

                        String evaluationResult = ModelApiClient.getAnswer(prompt, jsonData);
                        if (evaluationResult == null || evaluationResult.trim().isEmpty()) {
                            throw new RuntimeException("API返回的评测结果为空");
                        }

                        batchEvaluates.add(Evaluate.builder()
                                .singleReport(evaluationResult)
                                .active(true)
                                .answer(llmAnswer)
                                .judge(judge)
                                .build());
                    }

                    // 立即提交当前批次处理结果
                    evaluateRepository.saveAll(batchEvaluates);
                    System.out.println("已完成评测批次: " + (i/batchSize + 1) + "/" + (int)Math.ceil((double)llmAnswers.size()/batchSize));
                }

                System.out.println("完成裁判模型: " + judge);
            }
            return ResponseEntity.ok(Map.of(
                "success", true,
                "message", "评测完成"
            ));
        } catch (Exception e) {
            return ResponseEntity.internalServerError().body(Map.of(
                "success", false,
                "errorType", e.getClass().getSimpleName(),
                "message", "评测失败",
                "detail", e.getMessage(),
                "timestamp", System.currentTimeMillis()
            ));
        }
    }

    /**
     * 更安全的字符串处理方法
     * 1. 保留基本格式符号
     * 2. 仅处理真正危险的字符
     */
    private String sanitizeString(String input) {
        if (input == null) {
            return "";
        }
        // 仅处理真正会导致JSON解析问题的字符
        return input.replace("\\", "\\\\")  // 保留转义字符
                   .replace("\"", "\\\"")   // 转义引号
                   .replace("\n", "\\n")    // 保留换行但转义
                   .replace("\r", "\\r")    // 保留回车但转义
                   .replace("\t", "\\t");   // 保留制表符但转义
    }

    private String generateJsonData(StandardAnswer standardAnswer, StandardQuestion standardQuestion, LlmAnswer llmAnswer) {
        try {
            ObjectMapper objectMapper = new ObjectMapper();
            Map<String, Object> jsonMap = new LinkedHashMap<>();

            // 添加基础字段（经过安全处理）
            jsonMap.put("question", sanitizeString(standardQuestion.getQuestion()));
            jsonMap.put("reference_answer", sanitizeString(standardAnswer.getAnswer()));

            // 添加关键点（按参考样例格式优化）
            ResponseEntity<List<Points>> response = pointsService.getPointsByAnswerId(
                standardAnswer.getQuestionId(),
                standardAnswer.getVersion()
            );

            if (response != null && response.getStatusCode().is2xxSuccessful()) {
                List<Points> pointsList = response.getBody();
                if (!pointsList.isEmpty()) {
                    Map<String, String> keyPointsMap = new LinkedHashMap<>();
                    int index = 1;
                    for (Points point : pointsList) {
                        if (point.getActive()) {
                            keyPointsMap.put(
                                String.valueOf(index++),
                                sanitizeString(point.getKeyword())
                            );
                        }
                    }
                    if (!keyPointsMap.isEmpty()) {
                        jsonMap.put("key_points", keyPointsMap);
                    }
                }
            }

            // 最后添加候选答案（按参考样例顺序）
            jsonMap.put("candidate_answer", sanitizeString(llmAnswer.getAnswer()));

            // 序列化为格式化的JSON字符串
            return objectMapper.writerWithDefaultPrettyPrinter().writeValueAsString(jsonMap);
        } catch (JsonProcessingException e) {
            throw new RuntimeException("生成评测JSON数据失败: " + e.getMessage(), e);
        }
    }


}